from s4dd import S4forDenovoDesign
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# 路径
MODEL_PATH = "./new_models/s4_finetuned_final"  # 已训练的模型路径
TRAINING_DATA_PATH = "./datasets/chemblv31/train_30_filtered.txt"  # 用于微调的30%数据
MAX_LENGTH = 50  # 限制生成序列的最大长度
EPOCHS = 10  # 微调训练的轮数

def main():
    # -------------------------------
    # 1. 加载预训练的S4模型
    # -------------------------------
    s4 = S4forDenovoDesign.from_file(MODEL_PATH)

    # -------------------------------
    # 2. 微调训练
    # -------------------------------
    for epoch in range(EPOCHS):
        print(f"Starting epoch {epoch+1}...")

        # 进行微调训练
        s4.train(
            training_molecules_path=TRAINING_DATA_PATH,  # 使用筛选后的数据进行微调
            val_molecules_path="./datasets/chemblv31/valid.zip",  # 验证集路径
        )

        # 每轮保存微调后的模型
        save_path = f"./new_models/s4_length_constrained_finetuned_{epoch+1}"
        s4.save(save_path)
        print(f"Model saved at {save_path}")

    print("Fine-tuning complete, models saved for each epoch.")

if __name__ == '__main__':
    main()
